"""流式响应智能体"""

import asyncio
import time
from typing import AsyncGenerator, Dict, Any, Optional, List
from loguru import logger
import traceback

from agent.react_agent import ReactAgent
from agent.agent_state import AgentState
from model.event_model import EventType, StreamingEvent


class StreamingReactAgent(ReactAgent):
    """实时响应的流式智能体"""
    
    def __init__(self, verbose: bool = False, mode: str = "hybrid", enable_reflection: bool = True):
        super().__init__(verbose, mode, enable_reflection=enable_reflection)
        self._event_queue = None
        self._current_step = 0
    
    async def run_stream(self,
            session_id: str = None,
            query: str = None,
            interrupt_feedback: str = None,
            history: List[Any] = None,
            max_steps: int = None
    ) -> AsyncGenerator[StreamingEvent, None]:
        self._event_queue = asyncio.Queue()
        self._current_step = 0

        # 在后台任务中启动智能体执行
        task = asyncio.create_task(
            self._run_with_events(session_id, query, interrupt_feedback, history, max_steps)
        )
        try:
            while True:
                try:
                    event = await asyncio.wait_for(self._event_queue.get(), timeout=1.0)
                    logger.info(event)
                    await asyncio.sleep(0)
                    yield event

                    if event.type in [EventType.COMPLETE, EventType.ERROR]:
                        break
                        
                except asyncio.TimeoutError:
                    if task.done():
                        break
                    continue
        finally:
            if not task.done():
                task.cancel()
                try:
                    await task
                except asyncio.CancelledError:
                    logger.error(asyncio.CancelledError)
                    pass
    
    async def _run_with_events(
            self,
            session_id: str = None,
            query: str = None,
            interrupt_feedback: str = None,
            history: List[Any] = None,
            max_steps: int = None
    ):
        try:
            response = await self.run(session_id, query, interrupt_feedback, history, max_steps)
            await self._emit_event(EventType.COMPLETE, {
                "response": response,
                "success": response["success"]
            })

        except Exception as e:
            logger.error(traceback.format_exc())
            await self._emit_event(EventType.ERROR, {
                "error": str(e),
                "success": False
            })
    
    async def _emit_event(self, event_type: EventType, data: Dict[str, Any], metadata: Optional[Dict[str, Any]] = None):
        if self._event_queue:
            event = StreamingEvent(
                type=event_type,
                data=data,
                timestamp=time.time(),
                step=self._current_step,
                metadata=metadata or {}
            )
            await self._event_queue.put(event)

    async def _decide_approach_node(self, state: AgentState) -> AgentState:
        result_state = await super()._decide_approach_node(state)

        await self._emit_event(EventType.APPROACH,  {
                "chosen_approach": result_state['chosen_approach']
            }
        )
        return result_state

    async def _think_node(self, state: AgentState) -> AgentState:
        self._current_step = state['current_step'] + 1

        await self._emit_event(EventType.THINKING_START, {
            "step": self._current_step,
            "input": state["input"]
        })

        result_state = await super()._think_node(state)
        
        # 发送包含思考内容的事件
        if result_state["thoughts"]:
            latest_thought = result_state["thoughts"][-1]
            await self._emit_event(EventType.THINKING, {
                "thought": latest_thought,
                "step": self._current_step
            })
        
        # 若某项行动已规划，则发送行动已规划事件
        if result_state["actions"] and len(result_state["actions"]) > len(result_state["observations"]):
            latest_action = result_state["actions"][-1]
            await self._emit_event(EventType.ACTION_PLANNED, {
                "action": latest_action.get("name", latest_action.get("action", "unknown")),
                "input": latest_action.get("input", latest_action.get("action_input", {})),
                "step": self._current_step
            })
        
        return result_state
    
    async def _act_node(self, state: AgentState) -> AgentState:
        if not state["actions"]:
            return state
        
        current_action = state["actions"][-1]
        action_name = current_action.get("name", "unknown")
        action_input = current_action.get("input", "")

        await self._emit_event(EventType.ACTION_START, {
            "action": action_name,
            "input": action_input,
            "step": self._current_step
        })

        result_state = await super()._act_node(state)

        # 发送动作执行结果
        if result_state["tool_results"]:
            latest_result = result_state["tool_results"][-1]
            await self._emit_event(EventType.ACTION_RESULT, {
                "tool": latest_result.get("tool", "unknown"),
                "result": latest_result.get("result", {}),
                "step": self._current_step
            })
        
        return result_state
    
    async def _observe_node(self, state: AgentState) -> AgentState:
        result_state = await super()._observe_node(state)
        
        # 发送观察结果
        if result_state["observations"]:
            latest_observation = result_state["observations"][-1]
            await self._emit_event(EventType.OBSERVATION, {
                "observation": latest_observation,
                "step": self._current_step
            })
        
        return result_state

    async def _plan_node(self, state: AgentState) -> AgentState:
        await self._emit_event(EventType.PLAN_START, {
            "step": self._current_step
        })

        result_state = await super()._plan_node(state)

        if result_state.get("current_plan"):
            plan = result_state["current_plan"]
            await self._emit_event(EventType.PLAN_CREATED, {
                "plan": {
                    "description": getattr(plan, 'description', 'No description'),
                    "confidence": getattr(plan, 'confidence', 0.0),
                    "steps": [
                        {
                            "id": getattr(step, 'id', f"step_{i}"),
                            "description": getattr(step, 'description', 'No description'),
                            "tool": getattr(step, 'tool', 'unknown')
                        }
                        for i, step in enumerate(getattr(plan, 'steps', []))
                    ]
                },
                "step": self._current_step
            })
        
        return result_state
    
    async def _execute_node(self, state: AgentState) -> AgentState:
        if state.get("current_plan"):
            plan = state["current_plan"]

            steps = getattr(plan, 'steps', [])
            for i, step in enumerate(steps):
                await self._emit_event(EventType.PLAN_STEP_START, {
                    "step_id": getattr(step, 'id', f"step_{i}"),
                    "step_description": getattr(step, 'description', 'No description'),
                    "tool": getattr(step, 'tool', 'unknown'),
                    "plan_step": i + 1,
                    "total_steps": len(steps)
                })

        result_state = await super()._execute_node(state)
        
        # 根据执行结果发送计划完成事件
        if result_state.get("is_complete"):
            await self._emit_event(EventType.PLAN_STEP_COMPLETE, {
                "status": "completed",
                "step": self._current_step
            })
        elif result_state.get("plan_failed"):
            await self._emit_event(EventType.PLAN_STEP_COMPLETE, {
                "status": "failed",
                "step": self._current_step
            })
        
        return result_state

    async def _adaptive_replan_node(self, state: AgentState) -> AgentState:

        await self._emit_event(EventType.REPLAN_START, {
            "step": self._current_step,
        })

        result_state = await super()._adaptive_replan_node(state)

        if result_state.get("current_plan"):
            plan = result_state["current_plan"]
            await self._emit_event(EventType.PLAN_CREATED, {
                "plan": {
                    "description": getattr(plan, 'description', 'No description'),
                    "confidence": getattr(plan, 'confidence', 0.0),
                    "steps": [
                        {
                            "id": getattr(step, 'id', f"step_{i}"),
                            "description": getattr(step, 'description', 'No description'),
                            "tool": getattr(step, 'tool', 'unknown')
                        }
                        for i, step in enumerate(getattr(plan, 'steps', []))
                    ]
                },
                "step": self._current_step
            })
        return result_state

    async def _reflect_node(self, state: AgentState) -> AgentState:
        if not self.reflection_module:
            return state

        await self._emit_event(EventType.REFLECTION_START, {
            "original_response": state.get("output", ""),
            "quality_threshold": self.reflection_module.quality_threshold,
            "step": self._current_step
        })

        result_state = await super()._reflect_node(state)
        reflection_metadata = result_state.get("metadata", {}).get("reflection", {})
        
        if reflection_metadata:
            # 为每次迭代发送批判反思事件
            reflection_history = reflection_metadata.get("reflection_history", [])
            for iteration_data in reflection_history:
                critique = iteration_data.get("critique")
                if critique:
                    if hasattr(critique, '__dict__'):
                        critique_data = {
                            "overall_quality": critique.overall_quality,
                            "confidence": critique.confidence,
                            "issues_count": len(critique.issues),
                            "strengths_count": len(critique.strengths),
                            "needs_refinement": critique.needs_refinement,
                            "reasoning": critique.reasoning,
                            "issues": [
                                {
                                    "type": issue.type.value if hasattr(issue.type, 'value') else str(issue.type),
                                    "severity": issue.severity.value if hasattr(issue.severity, 'value') else str(issue.severity),
                                    "description": issue.description,
                                    "suggestion": issue.suggestion,
                                    "confidence": issue.confidence
                                } for issue in critique.issues
                            ],
                            "strengths": critique.strengths
                        }
                    else:
                        # Dictionary fallback
                        critique_data = {
                            "overall_quality": critique.get("overall_quality", 0.0),
                            "confidence": critique.get("confidence", 0.0),
                            "issues_count": len(critique.get("issues", [])),
                            "strengths_count": len(critique.get("strengths", [])),
                            "needs_refinement": critique.get("needs_refinement", False),
                            "reasoning": critique.get("reasoning", ""),
                            "issues": critique.get("issues", []),
                            "strengths": critique.get("strengths", [])
                        }
                    
                    await self._emit_event(EventType.REFLECTION_CRITIQUE, {
                        "iteration": iteration_data.get("iteration", 1),
                        "critique": critique_data,
                        "step": self._current_step
                    })
            
            # 若有改进，则发送优化事件
            improvements = reflection_metadata.get("total_improvements", [])
            if improvements:
                await self._emit_event(EventType.REFLECTION_REFINEMENT, {
                    "improvements": improvements,
                    "quality_improvement": reflection_metadata.get("final_quality_score", 0.0) - 0.5,
                    "original_response": state.get("original_response", ""),
                    "refined_response": result_state.get("output", ""),
                    "step": self._current_step
                })
        
        # 发送反思完成事件
        await self._emit_event(EventType.REFLECTION_COMPLETE, {
            "final_quality_score": reflection_metadata.get("final_quality_score", 0.0),
            "reflection_iterations": reflection_metadata.get("reflection_iterations", 0),
            "threshold_met": reflection_metadata.get("threshold_met", False),
            "total_improvements": len(reflection_metadata.get("total_improvements", [])),
            "step": self._current_step
        })
        
        return result_state

    async def _finish_node(self, state: AgentState) -> AgentState:
        final_state = await super()._finish_node(state)
        response = {
            "input": final_state.get('input', ""),
            "output": final_state.get("output", "No output generated") \
                if isinstance(final_state, dict) else "No output generated",
            "steps": self._format_steps(final_state),
            "success": not final_state.get("has_error", False) \
                if isinstance(final_state, dict) else False,
            "error": final_state.get("error_message") \
                if isinstance(final_state, dict) else f"Invalid state type: {type(final_state)}",
            "metadata": {
                **(final_state.get("metadata", {}) if isinstance(final_state, dict) else {}),
                "mode": final_state.get("mode") if isinstance(final_state, dict) else self.mode,
                "session_id": final_state.get("session_id"),
                "chosen_approach": final_state.get("chosen_approach"),
                "current_plan": final_state.get("current_plan") \
                    if isinstance(final_state, dict) else None,
                "execution_time": f'{1.0:.2f}s'
            }
        }
        await self._emit_event(EventType.FINISH, {
            "response": response,
            "success": response["success"]
        })
        return final_state


class StreamingChatbot:
    """支持流式响应的聊天机器人"""
    
    def __init__(self, verbose: bool = False, mode: str = "hybrid", enable_reflection: bool = True):
        self.agent = StreamingReactAgent(verbose=verbose, mode=mode, enable_reflection=enable_reflection)

    async def chat_stream(
            self,
            session_id: str = None,
            query: str = None,
            interrupt_feedback: str = None,
            history: List[Any] = None,
            max_steps: int = None
    )-> AsyncGenerator[StreamingEvent, None]:
        try:
            async for event in self.agent.run_stream(
                session_id=session_id,
                query=query,
                interrupt_feedback=interrupt_feedback,
                history=history,
                max_steps=max_steps
            ):
                await asyncio.sleep(0)
                yield event

        except Exception as e:
            logger.error(traceback.format_exc())
            yield StreamingEvent(
                type=EventType.ERROR,
                data={
                    "error": str(e),
                    "success": False
                },
                timestamp=time.time(),
                step=0
            )